import os
import tarfile

import cv2
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image, ImageColor
from rasterio.io import MemoryFile


def read_tar_files():
    parent_directory = 'E:\\map data\\G1WBM'
    directory = os.listdir(parent_directory)

    #0: Land
    keys_and_names = {0: 'Other data', 1: 'Land (No Landsat observation)', 10: 'Snow',
                      20: 'Wet Soil / Wet Vegetation / Lava', 30: 'Salt Marsh',
                      40: 'Temporal Flooded Area', 50: 'Permanent Water',
                      51: 'Permanent Water (Added by SWBD)', 99: 'Ocean (Given by external land/sea mask)'}

    keys_and_colors = {0: '#FED1E2', 1: '#00FF00', 10: '#E1FEA8', 20: '#774342', 30: '#EE3300',
                       40: '#FFD190', 41: '#FF9F0A', 50: '#8ABDFF', 51: '#38D0FF', 52: '#0B67E0',
                       99: '#116DFF'}

    keys = np.array(list(keys_and_colors.keys()))

    tiny_image_numpy = np.array([[21, 38, 599], [1999, 2000, 2001]])
    # tiny_image_numpy_cv2 = cv2.cvtColor(tiny_image_numpy, cv2.COLOR_GRAY2RGB)

    all_width = 21600; all_height = 10800

    rgb_tuple = ImageColor.getcolor(keys_and_colors[99], "RGB")  # 返回元组 (255, 0, 0)
    rgb_array = list(rgb_tuple)

    all_images = np.zeros((all_height, all_width, 3), dtype=np.uint8)
    all_images[:, :, 0] = rgb_array[0]
    all_images[:, :, 1] = rgb_array[1]
    all_images[:, :, 2] = rgb_array[2]

    single_width = int(all_width / (360 / 5)); single_height = int(all_height / (180 / 5))
    # all_images = all_images + 99
    show_time = 0
    current_row = -1; current_column = -1

    band_r = np.array([1, 1, 1], dtype=np.uint8); band_g = band_r.copy(); band_b = band_r.copy()

    for file in directory:
        directory_name = os.path.join(parent_directory, file)

        child_names = os.listdir(directory_name)
        for child_name in child_names:
            if child_name.endswith('.tar'):
                tar = tarfile.open(os.path.join(directory_name, child_name))
                tar_members = tar.getmembers()
                for tar_member in tar_members:
                    print(f'File: {file} {directory.index(file)} / {len(directory)}; Tile name: {tar_member.name}')
                    tar_member_name = tar_member.name
                    if tar_member.name.endswith('.tif'):
                        tar_member_name1 = tar_member_name.split('/')[1]
                        lon_lat_name0 = tar_member_name1.split('_')[0]

                        east_west = 'w'; south_north = 'n'
                        if lon_lat_name0.find('e') != -1: east_west = 'e'
                        if lon_lat_name0.find('s') != -1: south_north = 's'
                        lon_lat_split = lon_lat_name0.split(east_west)
                        lon_lat_split[0] = lon_lat_split[0].replace(south_north, '')
                        lon_start = int(lon_lat_split[1]); lat_start = int(lon_lat_split[0])
                        if east_west == 'w': lon_start = -abs(lon_start)
                        if south_north == 's': lat_start = -abs(lat_start)

                        lon_start_pixel = round((lon_start - (-180)) / (180 - (-180)) * all_width)
                        lat_start_pixel = round((1 - (lat_start - (-90)) / (90 - (-90))) * all_height) - single_height
                        extracted_tif_file = tar.extractfile(tar_member)
                        content = extracted_tif_file.read()

                        with MemoryFile(content) as memfile:
                            # 打开内存文件中的 TIFF 数据集
                            with memfile.open() as dataset:
                                # 示例：读取影像数据（numpy 数组）
                                image_nparray_0 = dataset.read()  # shape: (波段数, 高度, 宽度)
                                image_nparray = image_nparray_0[0]

                                # image_nparray_unique = np.unique(image_nparray)
                                print(image_nparray.shape)
                                current_row = image_nparray.shape[0]; current_column = image_nparray.shape[1]
                                if band_r.shape[0] != current_row or band_r.shape[1] != current_column:
                                    band_r = np.zeros((image_nparray.shape[0], image_nparray.shape[1]), dtype=np.uint8)
                                    band_r = band_r + 255
                                    band_g = band_r.copy(); band_b = band_r.copy()

                                for value in keys_and_colors:
                                    hex_color = keys_and_colors[value]
                                    rgb_tuple = ImageColor.getcolor(hex_color, "RGB")  # 返回元组 (255, 0, 0)
                                    rgb_array = list(rgb_tuple)

                                    band_r[image_nparray == value] = rgb_array[0]
                                    band_g[image_nparray == value] = rgb_array[1]
                                    band_b[image_nparray == value] = rgb_array[2]

                                keys_not_included = np.array([])
                                # keys_not_included = np.setdiff1d(image_nparray_unique, keys, assume_unique=True)
                                # print(f"Unique: {keys_not_included}")
                                for value in keys_not_included:
                                    band_r[image_nparray == value] = 255
                                    band_g[image_nparray == value] = 255
                                    band_b[image_nparray == value] = 0

                                # 示例：获取元信息（如投影、变换参数等）
                                bounds = dataset.bounds
                                meta = dataset.meta
                                crs = dataset.crs
                                transform = dataset.transform

                                lon_range = bounds.right - bounds.left
                                lat_range = bounds.bottom - bounds.top

                                t = 11

                                # print(np.unique(image_nparray))
                                # image_cv2 = cv2.cvtColor(image_nparray)
                                resized_image_cv2_r = cv2.resize(band_r, (single_width, single_height), interpolation=cv2.INTER_AREA)
                                resized_image_cv2_g = cv2.resize(band_g, (single_width, single_height), interpolation=cv2.INTER_AREA)
                                resized_image_cv2_b = cv2.resize(band_b, (single_width, single_height), interpolation=cv2.INTER_AREA)

                                resized_image_cv2 = np.stack((resized_image_cv2_r, resized_image_cv2_g, resized_image_cv2_b), axis=-1, dtype=np.uint8)
                                all_images[lat_start_pixel:(lat_start_pixel + single_height), lon_start_pixel:(lon_start_pixel + single_width)] = resized_image_cv2

                        extracted_tif_file.close()
                tar.close()


    # plt.imshow(all_images)
    # plt.show()

    image = Image.fromarray(all_images)
    image.show()
    # show_time += 1
    t = 12

    return 0

if __name__ == '__main__':
    read_tar_files()